
import pyocr
import importlib
import sys
import time
import re
import io

importlib.reload(sys)
sys.stdout = io.TextIOWrapper(sys.stdout.buffer,encoding='UTF-8')
time1 = time.time()
# print("初始时间为：",time1)

import os.path
from pdfminer.pdfparser import PDFParser
from pdfminer.pdfdocument  import PDFDocument
from pdfminer.pdfinterp import PDFResourceManager, PDFPageInterpreter
from pdfminer.converter import PDFPageAggregator
from pdfminer.layout import LTTextBoxHorizontal, LAParams
from pdfminer.layout import LAParams, LTTextBox, LTTextLine, LTFigure, LTImage
#from pdfminer.pdfinterp import PDFTextExtractionNotAllowed
from pdfminer.pdfpage import PDFTextExtractionNotAllowed,PDFPage
import glob

from translate import Transtor


from concurrent.futures import ThreadPoolExecutor,as_completed
pool = ThreadPoolExecutor(12)

class PaperParser():
    def __init__(self):
        self.section = ["Abstract"]
        self.layout = None

    def parser(self,filename):
        '''
        加载文档对象
        :param pdf_file_path:
        :return:
        '''
        read_pdf = open(filename, 'rb')  # 打开PDF文件。
        # parser_pdf = PDFParser(read_pdf)  # 用文件对象创建一个PDF文档分析器。
        # pdf_document = PDFDocument(parser_pdf)  # 创建一个PDF文档。
        #
        # parser_pdf.set_document(pdf_document)
        # pdf_document.set_parser(parser_pdf)  # 连接分析器 与文档对象。
        # pdf_document.initialize()  # 如果没有密码，就创建一个空的字符串。

        # if not pdf_document.is_extractable:  # 检测文档是否提供txt转换，不提供就忽略。
        #     raise PDFTextExtractionNotAllowed
        if 0:
            pass
        else:
            pdf_manager = PDFResourceManager()  # 创建PDF资源管理器 来管理共享资源。
            pdf_laparams = LAParams()  # 创建一个PDF参数分析器。
            pdf_device = PDFPageAggregator(pdf_manager, laparams=pdf_laparams)  # 创建一个聚合器
            pdf_interpreter = PDFPageInterpreter(pdf_manager, pdf_device)  # 创建一个PDF页面解释器对象

            res = {}
            first = True # 获取第一页第一行为标题

            # 循环遍历列表，每次处理一页的内容，pdf_document.get_pages()获取page列表
            for each_page in PDFPage.get_pages(read_pdf):
                pdf_interpreter.process_page(each_page)  # 使用页面解释器来读取
                layout = pdf_device.get_result()  # 这里layout是一个LTPage对象 里面存放着这个page解析出的各种对象 一般包括LTTexBox,LTFigure,LTImage,
                # LTTexBoxHorizontal等等 想要获取文本就获得对象的text属性。
                line = 0
                outputLine = -1
                for each_info in layout:
                    if isinstance(each_info, LTTextBoxHorizontal):
                        results = each_info.get_text()
                        if line == 0 and first:
                            first = False
                            res["title"] = results
                            print(results)
                        for section in self.section:
                            if results == section + "\n":
                                outputLine = line + 1
                                print()
                            if line == outputLine:
                                res[section] = results.replace('-\n','').replace('\n','')
                        # 获取文本对象
                    if isinstance(each_info, LTTextBox):
                        pass
                        # print(each_info.get_text().strip())
                        # 获取图片对象
                    if isinstance(each_info, LTImage):
                        pass#print('这里获取到一张图片')
                        # 获取 figure 对象
                    if isinstance(each_info, LTFigure):
                        pass#print('这里获取到一个 figure 对象')
                    # 获取文本对象
                    line = line + 1
        return res


    def getSection(self,layout,section="Abstract"):
        '''
        解析每个章节内容，如：Abstract，Reference
        :param section:
        :return:
        '''
        line = 0
        nextLineisAbstract = -1
        for each_info in layout:
            # 获取文本对象
            if isinstance(each_info, LTTextBox):
                pass
                # print(each_info.get_text().strip())
            # 获取图片对象
            if isinstance(each_info, LTImage):
                print('这里获取到一张图片')
            # 获取 figure 对象
            if isinstance(each_info, LTFigure):
                print('这里获取到一个 figure 对象')
            # 获取文本对象
            if isinstance(each_info, LTTextBoxHorizontal):
                results = each_info.get_text()
                if line == 0:
                    pass  # print("title:",results)
                for section in self.section:
                    if results == section + "\n":
                        output = line + 1
                    if line == output:
                        return section + "\n"+results
                line = line + 1

    def read_ref(pdf_file_path):
        '''
        提取参考文献
        :param pdf_file_path:
        :return:
        '''
        #sys.setdefaultencoding("utf-8")

        # Open a PDF file.
        fp = open(pdf_file_path, 'rb')
        # Create a PDF parser object associated with the file object.
        parser = PDFParser(fp)
        # Create a PDF document object that stores the document structure.
        # Supply the password for initialization.
        rsrcmgr = PDFResourceManager()
        laparams = LAParams()
        device = PDFPageAggregator(rsrcmgr, laparams=laparams)
        interpreter = PDFPageInterpreter(rsrcmgr, device)
        document = PDFDocument(parser)

        # Process each page contained in the document.
        text_content = []
        for page in PDFPage.create_pages(document):
            interpreter.process_page(page)
            layout = device.get_result()
            for lt_obj in layout:
                if isinstance(lt_obj, LTTextBox) or isinstance(lt_obj, LTTextLine):
                    text_content.append(lt_obj.get_text())
                else:
                    pass

        # text_content 中每一个元素存储了一行文字
        total_text = ''.join(text_content).replace("\n", "")
        print(total_text)
        # 从字符串中解析出参考文献
        file = open("E:\Administrator\Desktop\\ref.txt", "w",encoding='utf-8')
        p = re.compile('\[\d+\][^\[\]]*\d\.')
        m = p.findall(total_text)
        for i in m:
            # print i
            if i.startswith("["):
                file.write(str(i))
                file.write("\n")
        file.close()

def saveToTxt(res,trans,saveFile):
    info = res["title"] + "Abstract\n" + res["Abstract"] + '\n' + '翻译：\n'+ str(trans,"utf8") + '\n'
    info +=  '################################################\n\n'
    info = str(info)
    print(info)
    with open(saveFile, 'a+',encoding='utf-8') as f:
        f.write(info)
    print(res["title"]+ "wtited to saveFile ")

def getAbstractMuti(path):
    paperParser = PaperParser()
    transtor = Transtor()

    if os.path.isdir(path):
        files = glob.glob("E:/论文/语义分割/_Semantic_segmentation/*.pdf")
        print(files)
        pool.map(getAbstract,files)
    else:
        pool.submit(getAbstract,path)

def getAbstract(filename):
    paperParser = PaperParser()
    transtor = Transtor()
    res = paperParser.parser(filename)
    trans = transtor.baiduTranslate(res["Abstract"]).encode('utf-8')
    #print(trans)
    saveToTxt(res, trans, os.path.join(os.path.dirname(filename), "Abstract.txt"))

if __name__ == '__main__':
    # pdf_file_path = get_pdf_file()
    #handleDirectory("")
    getAbstractMuti('E:/论文/语义分割/_Semantic_segmentation')
    pool.shutdown(wait=True)
    #read_ref('E:\Administrator\Desktop\DIOU_loss.pdf')